# -*- coding: utf-8 -*-
# @Author  : longbhu
# @Time    : 2025/3/20 14:25
# @Function:
import xarray as xr
import numpy as np
import rasterio
from rasterio.transform import from_origin
import os


def nc_to_tif(input_nc_path, variable_name, output_tif_path=None, resample_freq='1YE', resample_method='sum'):
    """
    将 NetCDF 文件转换为 GeoTIFF 文件。

    参数:
    - input_nc_path: 输入的 NetCDF 文件路径。
    - variable_name: NetCDF 文件中的变量名。
    - output_tif_path: 输出的 GeoTIFF 文件路径。如果未指定，则自动生成。
    - resample_freq: 重采样频率，默认为 '1YE'（每年）。
    - resample_method: 重采样方法，默认为 'sum'。
    """
    # 读取 NetCDF 数据
    nc_dataset = xr.open_dataset(input_nc_path)

    # 提取指定变量的数据
    if variable_name not in nc_dataset.data_vars:
        raise ValueError(f"变量 '{variable_name}' 不存在于 NetCDF 文件中。")

    variable = nc_dataset[variable_name]

    # 根据指定的频率和方法进行重采样
    if resample_method == 'sum':
        resampled_variable = variable.resample(valid_time=resample_freq).sum(dim='valid_time')
    elif resample_method == 'mean':
        resampled_variable = variable.resample(valid_time=resample_freq).mean(dim='valid_time')
    elif resample_method == 'max':
        resampled_variable = variable.resample(valid_time=resample_freq).max(dim='valid_time')
    elif resample_method == 'min':
        resampled_variable = variable.resample(valid_time=resample_freq).min(dim='valid_time')
    else:
        raise ValueError(f"不支持的重采样方法: {resample_method}")

    # 获取变量数据并移除多余维度
    array = resampled_variable.values.squeeze()  # 移除所有大小为1的维度

    # 设置输出 GeoTIFF 的参数
    transform = from_origin(
        resampled_variable.longitude[0].item(), # west：栅格数据的西边界（最小经度）。
        resampled_variable.latitude[0].item(),  # north：栅格数据的北边界（最大纬度）。
        resampled_variable.longitude[1].item() - resampled_variable.longitude[0].item(), # xres：栅格数据的水平分辨率（像素宽度）。
        resampled_variable.latitude[0].item() - resampled_variable.latitude[1].item()   # yres：栅格数据的垂直分辨率（像素高度）。
    )

    # 设置 TIFF 文件的元数据
    metadata = {
        'driver': 'GTiff',
        'count': 1,
        'dtype': array.dtype,
        'width': array.shape[1],
        'height': array.shape[0],
        'crs': 'EPSG:4326',
        'transform': transform
    }

    # 如果未指定输出路径，则自动生成
    if output_tif_path is None:
        base_name = os.path.basename(input_nc_path).replace('.nc', '')
        output_tif_path = os.path.join(os.path.dirname(input_nc_path),
                                       f"{base_name}_{resample_freq}_{resample_method}.tif")

    # 写入 TIFF 文件
    with rasterio.open(output_tif_path, 'w', **metadata) as dst:
        dst.write(array, 1)  # 写入第一个波段

    print(f"GeoTIFF 文件已保存到: {output_tif_path}")
    return output_tif_path


# 示例调用
input_nc_path = r'G:\GEP_data\input\era5\2023_precip_hour.nc'
variable_name = 'tp'
output_tif_path = nc_to_tif(input_nc_path, variable_name, resample_freq='1YE', resample_method='max')
print(output_tif_path)